from PySide6.QtWidgets import QApplication, QWidget
from PySide6.QtCore import Signal, Qt, QPoint, QRect
from PySide6.QtGui import (QPalette, QBrush, QPainter, QPen, QColor, QPixmap,
                           QImage)
import numpy as np
import pytesseract
import cv2
import sys


class CaptureWidget(QWidget):

    # 自定义信号
    send_str_signal = Signal(str)  # 将结果传给主界面
    close_capture_signal = Signal()

    def __init__(self, lang):
        super().__init__()

        self._lang = lang

        self.setWindowTitle('Capturing')
        self.setWindowFlags(Qt.FramelessWindowHint  # 无边框
                            | Qt.WindowStaysOnTopHint  # 页面置于最顶层
                            | Qt.Dialog)  # dialog样式，无按钮
        self.setWindowState(Qt.WindowFullScreen)  # 全屏

        # 获取当前屏幕的图像
        self.screen = QApplication.primaryScreen().grabWindow(0)

        # 配置调色板，将当前屏幕的图像画在widget上
        palette = QPalette()
        palette.setBrush(self.backgroundRole(), QBrush(self.screen))
        self.setPalette(palette)

        # QApplication.setOverrideCursor(Qt.CrossCursor)  # 设置全局的鼠标样式
        self.setCursor(Qt.CrossCursor)  # 设置鼠标样式为十字
        self.start, self.end = QPoint(), QPoint()  # 初始化起始点和终点

    def keyPressEvent(self, event):
        if event.key() == Qt.Key_Escape:  # esc键退出
            self.close_capture_signal.emit()  # 关闭窗口

        return super().keyPressEvent(event)

    def paintEvent(self, event):
        if self.start == self.end:
            return super().paintEvent(event)

        # 设置绘图方式
        painter = QPainter(self)
        painter.setPen(QPen(QColor(255, 0, 0), 3))  # 画笔颜色，粗细
        painter.drawRect(QRect(self.start, self.end))
        return super().paintEvent(event)

    def mousePressEvent(self, event):
        self.start.setX(int(event.position().x()))
        self.start.setY(int(event.position().y()))
        self.end.setX(int(event.position().x()))
        self.end.setY(int(event.position().y()))
        self.update()
        return super().mousePressEvent(event)

    def mouseMoveEvent(self, event):
        self.end.setX(int(event.position().x()))
        self.end.setY(int(event.position().y()))
        self.update()
        return super().mousePressEvent(event)

    def mouseReleaseEvent(self, event):
        if self.start == self.end:
            return super().mouseReleaseEvent(event)

        # 排序，让 (x1, y1) 为左上角的点，(x2, y2) 为右下角的点
        x1, x2 = sorted((self.start.x(), self.end.x()))
        y1, y2 = sorted((self.start.y(), self.end.y()))

        image = self.grab(QRect(QPoint(x1, y1), QPoint(
            x2, y2))).toImage()  # 在 QWidget 上截图，并转为 QImage
        image = self.QImageToCvMat(image)  # QImage 转 numpy 数组
        ocr_str = self.img2str(image).replace(' ', '')  # 删除字符串空格

        # 将识别的结果发送出去
        self.send_str_signal.emit(ocr_str)

        # 处理完关闭当前widget
        self.close_capture_signal.emit()

    # QImage 转 numpy 数组
    def QImageToCvMat(self, incomingImage) -> np.array:
        incomingImage = incomingImage.convertToFormat(
            QImage.Format.Format_RGB32)
        width = incomingImage.width()
        height = incomingImage.height()

        ptr = incomingImage.constBits()
        np_array = np.array(ptr).reshape(height, width, 4)  #  Copies the data
        return np_array

    # 解析图片
    def img2str(self, img) -> str:
        return pytesseract.image_to_string(img, lang=self._lang,
                                           timeout=5)  # lang 语言选择为中文


def main():
    app = QApplication(sys.argv)
    mainw = CaptureWidget('chi_sim')
    mainw.show()
    sys.exit(app.exec_())


if __name__ == '__main__':
    main()
